from typing import Optional
from book import Book, Page, Content, ContentType
import pdfplumber
from .exceptions import PageOutOfRangeException
from utils import LOG

class PDFParser:
    def __init__(self):
        pass

    # 因为 parse_pdf 没有使用到实例属性或方法，因此可以将其改成一个静态方法。添加一个 @staticmethod 注解即可
    @staticmethod
    def parse_pdf(input_file_path: str, page_num: Optional[int] = None) -> Book:
        book = Book(input_file_path)

        with pdfplumber.open(input_file_path) as pdf:
            if page_num is not None and page_num > len(pdf.pages):
                raise PageOutOfRangeException(len(pdf.pages), page_num)

            LOG.info(f'开始解析 {input_file_path}')
            if page_num is None:
                pages2parse = pdf.pages
            else:
                pages2parse = pdf.pages[:page_num]

            for pdf_page in pages2parse:
                page = Page()

                # 保留原始的文本内容
                raw_text = pdf_page.extract_text()
                tables = pdf_page.extract_tables()
                images = pdf_page.images

                # 因为表格内容也会被extract_text解析出来
                # 所以需要从原始文本中移除每个单元格的内容
                for table_data in tables:
                    for row in table_data:
                        for cell in row:
                            raw_text = raw_text.replace(cell, '', 1)

                # 处理文本内容
                if raw_text:
                    LOG.debug(f'[raw_text]\n{raw_text}')
                    raw_text_lines = raw_text.splitlines()
                    cleaned_raw_text_lines = [line.strip() for line in raw_text_lines if line.strip()]
                    cleaned_raw_text = "\n".join(cleaned_raw_text_lines)

                    text_content = Content(ContentType.TEXT, cleaned_raw_text)
                    page.add_content(text_content)
                    LOG.info(f'[cleaned_raw_text]\n{cleaned_raw_text}')

                # 处理表格内容
                if tables:
                    for table in tables:
                        table_str = ''
                        for row in table:
                            row_str = '\t'.join(row)  # 使用 \t 连接每一行的元素
                            table_str += row_str + '\n'  # 每行末尾添加换行符 \n
                        table_content = Content(ContentType.TABLE, table_str)
                        page.add_content(table_content)
                        LOG.info(f'[table]\n{table_str}')

                # 处理图片内容
                if images:
                    for img in images:
                        bbox = (img["x0"], img["top"], img["x1"], img["bottom"])
                        cropped_page = pdf_page.crop(bbox)
                        img = cropped_page.to_image(resolution=1080)  # 保存图片是1080p的
                        img_content = Content(ContentType.IMAGE, img)
                        page.add_content(img_content)
                        LOG.info(f'[image]\n{img}')

                book.add_page(page)

        return book


if __name__ == '__main__':
    book = PDFParser.parse_pdf('/home/bigmath/hjh/mini-ai-translator/test/test.pdf')
    for page in book.pages:
        for content in page.contents:
            print(content)